from argparse import ArgumentParser
from struct import unpack, pack
import sys
import hashlib
from Crypto.Cipher import AES

#aes_key = b'2B63B478DC23D5692B63B478DC23D569'
#aes_key = b'2B63B478DC23D5692B63B478DC23D569'
aes_key = b'ylsuxfhy}w{mh{|k5nn\x86\x87}nmhmxujt|}'

#v9[0] = 0x5A5A3257;
#v9[1] = 0x66975412;
#v9[2] = 0x66975412;
#v9[3] = 0x5A5A3257;


# iv = pack('<I', 0x5A5A3257) + pack('<I', 0x66975412) + \
#    pack('<I', 0x66975412) + pack('<I', 0x5A5A3257)
iv = b'1A52A367CB12C458'

# b'\x57\x32\x5a\x5a\x12\x54\x97\x66\x12\x54\x97\x66\x57\x32\x5a\x5a'


def main():
    parser = ArgumentParser(
        description='Read security configuration from MediaTek seccfg partition')
    parser.add_argument('file')
    args = parser.parse_args()

    file = open(args.file, 'rb')
    if unpack('<I', file.read(4))[0] != 0x4d4d4d4d:
        print('Invalid input file')
        sys.exit(1)

    version = unpack('<I', file.read(4))[0]
    unk0 = unpack('<I', file.read(4))[0]
    lock_state = unpack('<I', file.read(4))[0]
    unk1 = unpack('<I', file.read(4))[0]
    unk2 = unpack('<I', file.read(4))[0]
    if unpack('<I', file.read(4))[0] != 0x45454545:
        print('Invalid input file')
        sys.exit(1)
    sha256_encrypted = file.read(32)
    file.close()

    print('SECCFGv%d device lock state is %d ' % (version, lock_state), end='')
    if lock_state == 1:
        print('(locked)')
    elif lock_state == 3:
        print('(unlocked)')
    else:
        print('(unknown)')

    print('Encrypted SHA256 sum is:', sha256_encrypted)

    m = hashlib.sha256()
    m.update(pack('<I', 0x4d4d4d4d))
    m.update(pack('<I', version))
    m.update(pack('<I', unk0))
    m.update(pack('<I', lock_state))
    m.update(pack('<I', unk1))
    m.update(pack('<I', unk2))
    m.update(pack('<I', 0x45454545))
    computed_sha256 = m.digest()
    print('computed SHA: ', computed_sha256)

    cipher = AES.new(aes_key, AES.MODE_CBC, iv)
    sha256_computed_encrypted = cipher.encrypt(computed_sha256)
    print('Encrypted SHA256 sum is:', sha256_computed_encrypted)


if __name__ == '__main__':
    main()
